#pragma once
#ifndef __KDTREE_HPP__
#define __KDTREE_HPP__

#include "ViewPoint.hpp"
#include <vector>

namespace PhotonMapping
{
    using namespace NRenderer;
    using namespace std;

    struct kdNode {
        ViewPoint v;
        kdNode* leftChild;
        kdNode* rightChild;
        int split_dim;
        Vec3 bd_max, bd_min;
    };

    inline RGB _max(Vec3 a, Vec3 b) {
        return Vec3{ max(a.x, b.x), max(a.y, b.y), max(a.z, b.z) };
    }

    inline RGB _min(Vec3 a, Vec3 b) {
        return Vec3{ min(a.x, b.x), min(a.y, b.y), min(a.z, b.z) };
    }

    class kdTree {
    public:
        kdNode* root;
        void buildKDTree(kdNode* &pNode, vector<ViewPoint>& list, int l = -1, int r = -1, int dim = 0) {
			if (l == -1 && r == -1)
				l = 0, r = list.size();
			if (l >= r)
				return;
			int mid = (l + r) >> 1;
            switch (dim) {
            case 0:
                nth_element(list.begin() + l, list.begin() + mid, list.begin() + r, ViewPoint::comparer<0>());
            case 1:
                nth_element(list.begin() + l, list.begin() + mid, list.begin() + r, ViewPoint::comparer<1>());
            case 2:
                nth_element(list.begin() + l, list.begin() + mid, list.begin() + r, ViewPoint::comparer<2>());
            }
			pNode = new kdNode();
            pNode->v = list[mid];
			pNode->leftChild = pNode->rightChild = NULL;
			pNode->split_dim = dim;
			pNode->bd_max = pNode->v.pos;
			pNode->bd_min = pNode->v.pos;
			buildKDTree(pNode->leftChild, list, l, mid, (dim + 1) % 3);
			if (pNode->leftChild) {
				pNode->bd_max = _max(pNode->bd_max, pNode->leftChild->bd_max);
				pNode->bd_min = _min(pNode->bd_min, pNode->leftChild->bd_min);
			}
			buildKDTree(pNode->rightChild, list, mid + 1, r, (dim + 1) % 3);
			if (pNode->rightChild) {
				pNode->bd_max = _max(pNode->bd_max, pNode->rightChild->bd_max);
				pNode->bd_min = _min(pNode->bd_min, pNode->rightChild->bd_min);
			}
        }
        void delKdTree(kdNode* p) {
            if (!p)return;
            delKdTree(p->leftChild);
            delKdTree(p->rightChild);
            delete p;
        }
        void _query(kdNode* pNode, Vec3 pos, float r, vector<ViewPoint*>& result) {
            float dx = 0, dy = 0, dz = 0;
            if (pos.x > pNode->bd_max.x || pos.x < pNode->bd_min.x)
                dx = min(abs(pos.x - pNode->bd_max.x), abs(pos.x - pNode->bd_min.x));
            if (pos.y > pNode->bd_max.y || pos.y < pNode->bd_min.y)
                dy = min(abs(pos.y - pNode->bd_max.y), abs(pos.y - pNode->bd_min.y));
            if (pos.z > pNode->bd_max.z || pos.z < pNode->bd_min.z)
                dz = min(abs(pos.z - pNode->bd_max.z), abs(pos.z - pNode->bd_min.z));

            if (dx * dx + dy * dy + dz * dz > r * r)
                return;

            if (glm::distance(pNode->v.pos,pos) <= r)
                result.push_back(&(pNode->v));
            if (pNode->leftChild)
                _query(pNode->leftChild, pos, r, result);
            if (pNode->rightChild)
                _query(pNode->rightChild, pos, r, result);
        }
    public:
        kdTree() {};
        kdTree(vector<ViewPoint>& list, int l = -1, int r = -1, int dim = 0) {
            buildKDTree(root, list, l, r, dim);
        }
        ~kdTree() {
            delKdTree(root);
        }
        void query(Vec3 pos, float r, vector<ViewPoint*> &result) {
            _query(root, pos, r, result);
        }
    };
}

#endif
